﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Net.Sockets;
using System.Security.Cryptography;
using Lm.Common.Serialization;
using Lm.Common.ExtensionMethods;

namespace Lm.Common.Net
{
    /// <summary>
    /// 提供加密的网络流
    /// </summary>
    public class EncryptedNetworkStream 
    {
        protected NetworkStream NetStream { get; set; }

        protected byte[] Key { get; set; }
        protected byte[] IV { get; set; }

        protected EncryptedNetworkStream()
        { }

        public EncryptedNetworkStream(NetworkStream stream)
        { this.NetStream = stream; }

        /// <summary>
        /// 此方法必须在所使用的网络流开始传输数据之前调用，
        /// 否则可能会引起数据异常
        /// </summary>
        public void WaitEncrypt()
        {
            var publickKeyBuf = this.NetStream.ReadBlock();
            var publicKey = new BinaryFormatter().Deserialize<RSAParameters>(publickKeyBuf);
            using (var rsa = new RSACryptoServiceProvider())
            {
                rsa.ImportParameters(publicKey);
                using (var aes = new AesCryptoServiceProvider())
                {
                    this.NetStream.WriteBlock(rsa.Encrypt(aes.Key, true));
                    this.NetStream.WriteBlock(rsa.Encrypt(aes.IV, true));
                    this.NetStream.Flush();
                    this.Key = aes.Key;
                    this.IV = aes.IV;
                }
            }
        }

        /// <summary>
        /// 此方法必须在所使用的网络流开始传输数据之前调用，
        /// 否则可能会引起数据异常
        /// </summary>
        public void StartEncrypt()
        {
            using (var rsa = new RSACryptoServiceProvider())
            {
                var publicKey = rsa.ExportParameters(false);
                var publickKeyBuf = new BinaryFormatter().Serialize(publicKey);
                this.NetStream.WriteBlock(publickKeyBuf);
                this.NetStream.Flush();
                var keyBuf = this.NetStream.ReadBlock();
                var ivBuf = this.NetStream.ReadBlock();
                this.Key = rsa.Decrypt(keyBuf, true);
                this.IV = rsa.Decrypt(ivBuf, true);
            }
        }

        /// <summary>
        /// 与ReadBloc配对使用
        /// 对流加密完成后才能调用此方法，否则可能会造成数据异常。
        /// 加密流使用WaitEncrypt方法和StartEncrypt方法。
        /// </summary>
        /// <param name="buf"></param>
        public void WriteBlock(byte[] buf)
        {
            var data = new Lm.Common.Security.Cryptography.AesCrypto(this.Key, this.IV).Encrypt(buf);
            this.NetStream.WriteBlock(data);
        }

        /// <summary>
        /// 与ReadStringBlock配对使用
        /// 对流加密完成后才能调用此方法，否则可能会造成数据异常。
        /// 加密流使用WaitEncrypt方法和StartEncrypt方法。
        /// </summary>
        /// <param name="str"></param>
        public void WriteStringBlock(string str)
        {
            var data = new Lm.Common.Security.Cryptography.AesCrypto(this.Key, this.IV).Encrypt(Encoding.UTF8.GetBytes(str));
            this.NetStream.WriteBlock(data);
        }

        /// <summary>
        /// 与WriteBlock配对使用
        /// 对流加密完成后才能调用此方法，否则可能会造成数据异常。
        /// 加密流使用WaitEncrypt方法和StartEncrypt方法。
        /// </summary>
        /// <returns></returns>
        public byte[] ReadBlock()
        {
            var buf = this.NetStream.ReadBlock();
            return new Lm.Common.Security.Cryptography.AesCrypto(this.Key, this.IV).Decrypt(buf);
        }

        /// <summary>
        /// 与WriteStringBlock配对使用
        /// 对流加密完成后才能调用此方法，否则可能会造成数据异常。
        /// 加密流使用WaitEncrypt方法和StartEncrypt方法。
        /// </summary>
        /// <returns></returns>
        public string ReadStringBlock()
        {
            var buf = this.NetStream.ReadBlock();
            var data = new Lm.Common.Security.Cryptography.AesCrypto(this.Key, this.IV).Decrypt(buf);
            return Encoding.UTF8.GetString(data);
        }
    }
}
